#!/usr/bin/env python
#  Copyright (c) 2010
#  The Regents of the University of Michigan
#  All Rights Reserved

#  Permission is granted to use, copy, create derivative works, and
#  redistribute this software and such derivative works for any purpose,
#  so long as the name of the University of Michigan is not used in
#  any advertising or publicity pertaining to the use or distribution
#  of this software without specific, written prior authorization. If
#  the above copyright notice or any other identification of the
#  University of Michigan is included in any copy of any portion of
#  this software, then the disclaimer below must also be included.

#  This software is provided as is, without representation or warranty
#  of any kind either express or implied, including without limitation
#  the implied warranties of merchantability, fitness for a particular
#  purpose, or noninfringement.  The Regents of the University of
#  Michigan shall not be liable for any damages, including special,
#  indirect, incidental, or consequential damages, with respect to any
#  claim arising out of or in connection with the use of the software,
#  even if it has been or is hereafter advised of the possibility of
#  such damages.

from pydht import dht, dhtutil
import providencehelpers

import event

import array
import collections
import gzip
import itertools
import logging
import multiprocessing
import os
import Queue
import random
import signal
import socket
import SocketServer
import sys
import time
import types

TO_QUERY = -1
QUERIED = 0
DONE = 16 # 20

NUM_SENDERS = 3
# Cached to avoid dots in hot code.
compute_node_id = dht.compute_node_id
compute_shared_prefix_bits = dhtutil.compute_shared_prefix_bits

def on_ec2():
  return False  # TODO: stub

class ProvidenceError(Exception):
  pass

class ContactState(object):
  __slots__ = ['state']
               #'learn_count']
  def __init__(self):
    self.state = QUERIED
#    self.learn_count = 1


class Providence(object):
  '''The "Eye of Providence", or all-seeing eye.
  Scans Vuze routing tables.
  in_progress -- map address -> state.
  done -- map address we satisfactorily finished -> state where we quit.
  '''

  def __init__(self, external_ip, finished_callback, quit_after_scan=False):

    self.ip = external_ip
    self.clients = [Providence.Client(self, (external_ip, random.randint(1024, 65535)))]
    high_id_byte = ord(self.clients[0].dht_id[0]) ^ 0x80
    junk, port2 = dht.brute_force_id_with_bits((external_ip,),
                                                     chr(high_id_byte),
                                                     1)
    assert port2 is not None
    self.clients.append(Providence.Client(self, (external_ip, port2)))

    self.finished_callback = finished_callback
    self.quit_after_scan = quit_after_scan

  class Client(dht.DHTClient):
    '''Forwards DHT events to the owning Providence object.
    Used as a "dumb", "shallow" client that only does find-nodes.
    '''

    def __init__(self, parent, *args):
      dht.DHTClient.__init__(self, *args, firewalled=True)
      self.handle_find_node_reply = parent.handle_find_node_reply

  def _send_loop(self):
    clients = self.clients
    event.init()
    dhtutil.HACK_DISABLE_TIMER = False
    for client in clients:
      client._handle_recv = lambda self: None
    print 'worker, PID', os.getpid()
    event.signal(signal.SIGTERM, event.abort)
    event.signal(signal.SIGINT, event.abort)

    dhtutil.the_send_queue.send_slowed_callback = event.abort

    def handle_read_ready():
      q = self.to_contact_queue
      try:
        contacts = q.get(False)
        for contact in contacts:
          self._scan_contact(contact)
      except Queue.Empty:
        pass
      return True

    def start_debug_repl():
      from SocketConsole import SocketConsole
      class InteractiveConsoleHandler(SocketServer.BaseRequestHandler):
        def setup(self):
          self.data = ''
        def handle(self):
          SocketConsole(self.request).interact()
      class ReusingTCPServer(SocketServer.TCPServer):
        allow_reuse_address = True
      server = ReusingTCPServer(('localhost', 31337), InteractiveConsoleHandler)
      server.serve_forever()
          
    event.read(self.to_contact_queue._reader.fileno(), handle_read_ready)
    event.dispatch()
    print 'worker terminating'

  def _timer_event(self):
    if not any(x.is_alive() for x in self.senders):
      self.finished_callback(self, self.quit_after_scan)
      if self.quit_after_scan:
        event.abort()
      return None
    else:
      return True
      

  def handle_find_node_reply(self, reply, sender_contact, sender_address):
    '''Update the state machine for the returned contacts and the sender.
    sender is the TRUE sending address, not necessarily the one from the
    routing tables.'''

    contact_states = self.contact_states

    #logging.debug('got find node reply from %s:%s' % sender_address)
    # Mark responsive.
    contact_states.set_state(sender_address[0], sender_address[1], 1)

    # Take care of all the new contacts: if they're really new, scan 'em.
    to_query = []
    for contact in reply.contacts:
      # Faster version for when learn_count is disabled.
      addr = dht.get_contact_address(contact)
      if not contact_states.has_item(*addr):
        contact_states.set_state(addr[0], addr[1], 0)
        to_query.append(contact)
      # contact_state = get(contact)
      # if contact_state is None:
      #   contact_states[contact] = ContactState()
      #   _scan_contact(contact, 0)
      # else:
      #   contact_state.learn_count += 1

    if to_query:
      self.to_contact_queue.put(to_query)

  def _scan_contact(self, contact):
    # compute_node_id is cached dht.compute_node_id.
    # compute_shared_prefix_bits is cached from dhtutil.
    contact_id = compute_node_id(contact)

    # Pick whichever of the two clients differs in the first bit.
    clients = self.clients
    client = clients[(ord(clients[1].dht_id[0]) & ord(contact_id[0])) >= 128]

    contact_id = array.array('B', contact_id)
    send_find_node = client.send_find_node
    for bucket in xrange(0, DONE+1, 2):
      idx = bucket / 8
      flippy = (1 << (bucket %8))
      contact_id[idx] ^= flippy
      send_find_node(contact,contact_id.tostring())
      contact_id[idx] ^= flippy
    # providencehelpers.scan_contact_inner_loop(contact, contact_id,
    #                                           client.send_find_node)


  def start_scan(self, bootstrap_nodes=None):
    dhtutil.the_send_queue.send_slowed_callback = lambda: None

    event.timeout(30, self._timer_event)

    if not bootstrap_nodes:
      bootstrap_nodes = [(socket.gethostbyname('dht.aelitis.com'), 6881)]

#    self.pending = set()
#    self.in_progress = {}
#    self.done = set()

    # Number of times we learned about each node.
#    self.learn_count = collections.defaultdict(int)

    self.contact_states = providencehelpers.AddrPortDict()
    self.bogus_replies = 0

    self.to_contact_queue = multiprocessing.Queue()
    self.senders = []
    for _ in xrange(NUM_SENDERS):
      s = multiprocessing.Process(target=self._send_loop)
      self.senders.append(s)
      s.start()
      
    client = self.clients[0]
    for idx, node in enumerate(bootstrap_nodes):
      contact = dht.make_contact(node)
      self.contact_states.set_state(node[0], node[1], 0)
      client.send_find_node(contact, client.dht_id)

    assert dhtutil.HACK_DISABLE_TIMER

  def terminate_scan_early(self):
    for p in self.senders:
      p.terminate()

Providence.handle_find_node_reply = types.MethodType(
  providencehelpers.handle_find_node_reply, None, Providence)

def providence_finished(providence, quitting=False):
  the_time = time.time()
  outfyl = gzip.GzipFile(
      ('/mnt/' if on_ec2() else '') + 'scan.%s.%s.out.gz' % (dht.NAME, time.strftime('%Y-%m-%d.%H:%M:%S',
                                                     time.localtime(the_time))), 'w')
  
  def print_node_set(s):
    def myPred(x):
      return pred(x[1])
    print >> outfyl, 'List:'
    for contact in s:
      ip, port = dht.get_contact_address(contact)
      print >> outfyl, '%s:%d' % (ip, port)

  print >> outfyl, 'Crawl finished at %s!' % time.ctime(the_time)
  print >> outfyl, 'Bandwidth stats:'

  print >> outfyl, ('Bogus replies: %d' % providence.bogus_replies)

  print >> outfyl, ('Known nodes: %d' % len(providence.contact_states))
  providence.contact_states.print_known_nodes(outfyl)
  print >> outfyl, ('Responsive nodes: %d' % providence.contact_states.count_responsive_nodes())
  providence.contact_states.print_responsive_nodes(outfyl)

  # print >> outfyl, ('Fully-responsive nodes: %d' %
  #                   sum(1 for (c,s) in providence.contact_states.iteritems()
  #                         if s.state == DONE))
  # print_node_set(providence.contact_states, lambda state: state.state == DONE)

  # print >> outfyl, ('\nPartially-responsive nodes: %d' %
  #                   sum(1 for (c,s) in providence.contact_states.iteritems()
  #                       if s.state != 0 and s.state != DONE))
  # print_node_set(providence.contact_states,
  #                lambda s: s.state != 0 and s.state != DONE)

  # print >> outfyl, ('\nUnresponsive nodes: %d' % 
  #                   sum(1 for (c,s) in providence.contact_states.iteritems()
  #                       if s.state == 0))
  # print_node_set(providence.contact_states,
  #                lambda s: s.state == 0)

  # Disable learn_count.
  # print >> outfyl, '\nTimes we learned about each node:'
  # for contact, state in providence.contact_states.iteritems():
  #   ip, port = dht.get_contact_address(contact)
  #   print >> outfyl, '%s:%d\t%d\n' % (ip, port, state.learn_count)

  outfyl.close()
  if not quitting:
    print 'Starting another scan at %s!' % time.ctime()
    print 'PID:', os.getpid()
    providence.start_scan(bootstrap_nodes)



if __name__ == '__main__':
  logging.basicConfig(filename="providence.log",
                      level=logging.INFO)
                      #level=logging.DEBUG)

  only_once = len(sys.argv) > 2 and sys.argv[2].lower() == 'once'
  providence = Providence(sys.argv[1], providence_finished, only_once)

  bootstrap_nodes = []
  for idx in xrange(3 if only_once else 2, len(sys.argv) - 1, 2):
    addr = sys.argv[idx], int(sys.argv[idx + 1])
    bootstrap_nodes.append(addr)

  def quit_early():
    print 'quitting early'
    providence.terminate_scan_early()
    providence_finished(providence, quitting=True)
    event.abort()

  event.signal(signal.SIGTERM, quit_early)
  event.signal(signal.SIGINT, quit_early)

  print 'Starting at %s' % time.ctime()
  providence.start_scan(bootstrap_nodes)
  event.dispatch()
